"""
词汇表管理器 (Vocabulary Manager)
就像一本字典，负责管理AI认识的所有词汇
把文字和数字之间进行相互转换
"""

from typing import Dict, List, Optional

class Vocab:
    """
    词汇表类 - AI的字典
    
    主要功能：
    1. 给每个词分配一个唯一的数字ID（就像身份证号）
    2. 可以通过词找到ID，也可以通过ID找到词
    3. 管理AI认识的所有词汇
    
    例如：
    "你" ↔ 0
    "好" ↔ 1  
    "吗" ↔ 2
    """
    
    def __init__(self, words: List[str]):
        """
        初始化词汇表
        words: 词汇列表，比如 ["你", "好", "吗", "？"]
        """
        # 创建两个字典用于双向查找
        self.encode: Dict[str, int] = {}    # 词 → ID的映射，比如 "你" → 0
        self.decode: Dict[int, str] = {}    # ID → 词的映射，比如 0 → "你"
        self.words: List[str] = words       # 保存所有词汇的列表
        
        # 建立双向映射关系
        for i, word in enumerate(words):
            self.encode[word] = i    # 给每个词分配ID
            self.decode[i] = word    # 记录每个ID对应的词
    
    def encode_word(self, word: str) -> Optional[int]:
        """
        把词转换成ID（编码）
        
        输入：词汇字符串，比如 "你"
        输出：对应的ID数字，比如 0；如果词不存在则返回None
        
        就像在字典里查找某个字的页码
        """
        return self.encode.get(word)
    
    def decode_word(self, token_id: int) -> Optional[str]:
        """
        把ID转换成词（解码）
        
        输入：ID数字，比如 0
        输出：对应的词汇，比如 "你"；如果ID不存在则返回None
        
        就像根据页码找到对应的字
        """
        return self.decode.get(token_id)
    
    @staticmethod
    def default_words() -> List[str]:
        """
        默认词汇表 - 提供一个基础的词汇集合
        
        返回：包含常用词的列表
        这些词可以用来测试系统是否正常工作
        """
        return ["hello", "world", "this", "is", "rust", "</s>"]
    
    @classmethod
    def default(cls) -> 'Vocab':
        """
        创建默认词汇表实例
        
        返回：使用默认词汇的Vocab对象
        方便快速创建一个可用的词汇表进行测试
        """
        return cls(cls.default_words())

# 使用示例：
# vocab = Vocab(["你", "好", "吗", "？"])
# 
# # 编码：把词转换成数字
# word_id = vocab.encode_word("你")  # 返回 0
# 
# # 解码：把数字转换成词  
# word = vocab.decode_word(0)       # 返回 "你"
#
# # 检查词汇表大小
# vocab_size = len(vocab.words)     # 返回 4